#ifndef GMM__H
#define GMM__H

#include <iostream>
#include <Eigen/Core>
#include <Eigen/Eigenvalues> 
#include <random>
#include <pcl/common/transforms.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <ros/ros.h>

namespace gvm
{


struct normal_random_variable
{
    normal_random_variable(Eigen::MatrixXd const& covar)
        : normal_random_variable(Eigen::VectorXd::Zero(covar.rows()), covar)
    {}

    normal_random_variable(Eigen::VectorXd const& mean, Eigen::MatrixXd const& covar)
        : mean(mean)
    {
        Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eigenSolver(covar);
        transform = eigenSolver.eigenvectors() * eigenSolver.eigenvalues().cwiseSqrt().asDiagonal();
    }

    Eigen::VectorXd mean;
    Eigen::MatrixXd transform;

    Eigen::VectorXd operator()() const
    {
        static std::mt19937 gen{ std::random_device{}() };
        static std::normal_distribution<> dist;

        return mean + transform * Eigen::VectorXd{ mean.size() }.unaryExpr([&](double x) { return dist(gen); });
    }
};



class GMM
{
public:
	GMM(){}
	~GMM(){}
	GMM(std::vector<double> pi_vector,std::vector<Eigen::Vector3d> means, std::vector<Eigen::Matrix3d> var,int n_component,
		Eigen::Vector3d sensor_location):
		pi_vector_(pi_vector),
		means_(means),
		var_(var),
		n_component_(n_component),
		sensor_location_(sensor_location)
	{
		for(int i=0;i<n_component_;i++)
		{
			normal_random_variable random_variable(means[i].cast<double>(),var[i].cast<double>());
			randomGenerator.push_back(random_variable);
		}
	}

	std::vector<double> pi_vector_;
	std::vector<Eigen::Vector3d> means_;
	std::vector<Eigen::Matrix3d> var_;
	int n_component_;
	Eigen::Vector3d sensor_location_;
	int agent_id{0};
	int color_id{37};

	std::vector<normal_random_variable> randomGenerator;

public:

	inline double getMahalanobisDistance(Eigen::Vector3d point,Eigen::Vector3d means,Eigen::Matrix3d var)
	{
		double distance=(point-means).transpose()*var.inverse()*(point-means);
		return distance;
	}

	std::vector<Eigen::Vector3d> GMMRandomSamplePoints(int n_SamplePoints,double n_sigma)
	{
		std::vector<Eigen::Vector3d> samplePoints;
		for(int i=0;i<n_component_;i++)
		{
			int clusterPoint_num=n_SamplePoints*pi_vector_[i];
			for(int j=0;j<clusterPoint_num;j++)
			{
				Eigen::Vector3d tmp_point=randomGenerator[i]();

				if(getMahalanobisDistance(tmp_point,means_[i],var_[i])>pow(n_sigma,2))
					continue;
				samplePoints.push_back(tmp_point);
			}

		}

		return samplePoints;
	}

	void init(std::vector<double> pi_vector,std::vector<Eigen::Vector3d> means, std::vector<Eigen::Matrix3d> var,int n_component,
		Eigen::Vector3d sensor_location)
	{
		pi_vector_=pi_vector;
		means_=means;
		var_=var;
		n_component_=n_component;
		sensor_location_=sensor_location;
		randomGenerator.clear();
		for(int i=0;i<n_component_;i++)
		{
			normal_random_variable random_variable(means[i].cast<double>(),var[i].cast<double>());
			randomGenerator.push_back(random_variable);
		}
	}

	void setAgentId(int agentSetId, int colorSetId)
	{
		agent_id=agentSetId;
		color_id=colorSetId;
	}

	bool pointUnderFOV()
	{
		return true;

	}

	double getGaussianProbability(Eigen::VectorXd x,Eigen::VectorXd u,Eigen::MatrixXd var)
	{
		double PI=3.1415926;
		double tmp=1.0/pow(2*PI,1.5)/pow(var.determinant(),0.5)*exp(-0.5*(x-u).transpose()*var.inverse()*(x-u));
		// ROS_INFO_STREAM("GAUSSIAN u:"<<u<<" var: "<<var<<"var determinant: "<<var.determinant());
		// ROS_INFO("GAUSSIAN: %f",tmp);
		return tmp;
	}


	bool fit(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud,std::vector<Eigen::Vector3d>& means,std::vector<Eigen::Matrix3d>& var,int cluster_num, std::vector<double>& pi_vector,
			bool pointAtedge, bool allpointAtedge,int component_point_num=50)
	{
		double regularvalue=0.001;
		int size=cloud->points.size();

		int label[size];

		//kmeans
		Eigen::Vector3d centre[cluster_num];
		double label_count[cluster_num];

		std::vector<Eigen::Vector3d> last_means(cluster_num);
		int count=0;
		int new_cluster_num=cluster_num;

		while(true)
		{
			cluster_num=means.size();
			last_means=means;
			
			//INIT
			for(int i=0;i<cluster_num;i++)
			{
				centre[i]<<0,0,0;
				label_count[i]=0;
				// last_means[i]=means[i];
			}

			// Find the closest menas of clusters to each point in cloud
			for(int i=0;i<size;i++)
			{
				Eigen::Vector3d point(cloud->points.at(i).x,cloud->points.at(i).y,cloud->points.at(i).z);
				double tmp_dis=9999;
				for(int j=0;j<cluster_num;j++)
				{
					if((point-means[j]).norm()<=tmp_dis)
					{
						label[i]=j;
						tmp_dis=(point-means[j]).norm();
					}

				}
				centre[label[i]]+=point;
				label_count[label[i]]++;
			}

			new_cluster_num=cluster_num;
			std::vector<Eigen::Vector3d> tmp_means;
			
			for(int i=0;i<new_cluster_num;i++)
			{
				if(label_count[i]==0)
				{
					cluster_num--;
					// auto iter = tmp_means.erase(tmp_means.begin() + i);
				}
				else
					tmp_means.push_back(centre[i]/(double)(label_count[i]));
			}

			means=tmp_means;

			bool converge{true};

			for(int i=0;i<cluster_num;i++)
			{
				if(means.size()!=last_means.size()) // TODO
				{
					converge=false; 
					break;
				}

				if((means[i]-last_means[i]).norm()>0.00001)
				{
					converge=false;
					break;
				}
			}
			if(converge) break;

			count++;
			if(count>100)
			{
				ROS_WARN("KMEANS FAILED!!!!!");
				break;
			}
		}
		// std::cout << "count ==============================" << count << std::endl;


		for(int i=0;i<cluster_num;i++)
		{
			if(means[i](0)==NAN)
			{
				ROS_WARN("KMEANS ERROR");
			}
			// ROS_INFO("\033[1;%dm kmeans CENTRE: %f %f %f label_count: %f \033[0m",color_id,means[i](0),means[i](1),means[i](2),label_count[i]);
		}

		// ROS_INFO("\033[1;%dm kmeans use time %f cluster number:%d \033[0m",color_id,(end-start).toSec(),cluster_num);

		var.clear();
		for (int i = 0; i < cluster_num; i++)
		{
			var.push_back(Eigen::Matrix3d::	Zero());
		}

		std::vector<double> pi(cluster_num,0);

		for(int i=0;i<size;i++)
		{
			Eigen::Vector3d point(cloud->points.at(i).x,cloud->points.at(i).y,cloud->points.at(i).z);
			var[label[i]]+=(point-means[label[i]])*(point-means[label[i]]).transpose();
		}
		for(int i=0;i<cluster_num;i++)
		{
			var[i]=var[i]/(double)label_count[i];

			// ROS_INFO("label_count %d : %f ",i,label_count[i]);
			// ROS_INFO_STREAM("var: %f "<<var[i]);

			var[i](0,0)+=regularvalue;
			var[i](1,1)+=regularvalue;
			var[i](2,2)+=regularvalue;

			if(var[i].determinant()<=1e-10)
			{
				ROS_WARN("Var==0, GMM Failed!!!!!!!!");
				return false;
			}
			// ROS_INFO_STREAM("var[i]:"<<i<<" "<<var[i]);
			// std::cout<<"var[i]:"<<i<<" "<<var[i]<<std::endl;
			pi[i]=(double)label_count[i]/(double)size;
			// ROS_INFO("pi: %f ",pi[i]);
		}

		std::vector<Eigen::Vector3d> new_means=means;
		std::vector<Eigen::Vector3d> new_means_sum=means;
		std::vector<Eigen::Matrix3d> new_var_sum=var;
		// std::cout << "new_means.size()==============================" << new_means.size() << std::endl;
		// std::cout << "new_var_sum.size()==============================" << new_var_sum.size() << std::endl;
		// std::cout << "cluster_num)==============================" << cluster_num << std::endl;

		double gama[size][cluster_num];
		double curlog=121.1;
		double lastlog=564.2;

		int iter=0;
		while(true)
		{
			for(int k=0;k<cluster_num;k++)
			{
				label_count[k]=0;
				new_means[k]<<0,0,0;
				new_means_sum[k]<<0,0,0;
				new_var_sum[k]=Eigen::Matrix3d::Zero();
			}

			lastlog=curlog;
			curlog=0;
			std::vector<double> sumProbability(size,0);

			for(int i=0;i<size;i++)
			{
				Eigen::Vector3d point(cloud->points.at(i).x,cloud->points.at(i).y,cloud->points.at(i).z);
				//计算某点在各高斯分布下之和
				for(int k=0;k<cluster_num;k++)
				{
					sumProbability[i]+=pi[k]*getGaussianProbability(point,means[k],var[k]);
				}

				for(int k=0;k<cluster_num;k++)
				{
					gama[i][k]=pi[k]*getGaussianProbability(point,means[k],var[k])/sumProbability[i];
					label_count[k]+=gama[i][k];
					new_means_sum[k]+=gama[i][k]*point;
				}
			}

			for(int i=0;i<size;i++)
			{
				curlog+=log(sumProbability[i]);
			}
			
			if(fabs(curlog-lastlog)<0.01)
				break;
			// ROS_INFO("label_count : %f %f %f",label_count[0],label_count[1],label_count[2]);

			for(int k=0;k<cluster_num;k++)
			{
				means[k]=new_means_sum[k]/(double)label_count[k];
				for(int i=0;i<size;i++)
				{
					Eigen::Vector3d point(cloud->points.at(i).x,cloud->points.at(i).y,cloud->points.at(i).z);

					new_var_sum[k]+=gama[i][k]*(point-means[k])*(point-means[k]).transpose();

				}
				var[k]=new_var_sum[k]/(double)label_count[k];
				pi[k]=label_count[k]/(double)size;
				var[k](0,0)+=regularvalue;
				var[k](1,1)+=regularvalue;
				var[k](2,2)+=regularvalue;

				if(var[k].determinant()<=1e-10)
				{
					ROS_WARN("\033[1;%dm Var==0, GMM Failed!!!!!!!!  \033[0m",color_id);
					return false;
				}
			}
			iter++;
			if(iter>15)
			{
				ROS_WARN("\033[1;%dm GMM use to many iter, GMM FAILED !!!!!!! \033[0m",color_id);
				break;
			}
			// break;
		}

		pi_vector=pi;
		for(int i=0;i<cluster_num;i++)
		{
			if(!pointAtedge)
			{
				Eigen::EigenSolver<Eigen::Matrix3d> eigen_solver ( var[i] );
				Eigen::MatrixXd eig_value = eigen_solver.pseudoEigenvalueMatrix();
				Eigen::Matrix3d vector_martrix = eigen_solver.pseudoEigenvectors();
				//ROS_INFO("\033[1;%dm eigvalue: %f %f %f \033[0m",color_id,eig_value(0,0),eig_value(1,1),eig_value(2,2));
				//std::cout<<"vector_matrix:"<<vector_martrix<<std::endl;

				for(int j=0;j<3;j++)
				{
					if(eig_value(j,j)>pow(1.0,2))
					{
						Eigen::Vector3d tmp_vector=vector_martrix.col(j);
						if(fabs(tmp_vector(2))<0.3){
							ROS_INFO("\033[1;%dm fabs(tmp_vector(2))<0.3 \033[0m",color_id);
							return false;
						}
					}
				}
			}
			// ROS_INFO("\033[1;%dm GMM CENTRE: %f %f %f var determinant: %f \033[0m",color_id,means[i](0),means[i](1),means[i](2),var[i].determinant());
			// ROS_INFO_STREAM("var[i]:"<<i<<" "<<var[i]);

		}
		// ROS_INFO("\033[1;%dm GMM use time %f iter time: %d \033[0m",color_id,(end-start).toSec(),iter);
		return true;
	}



};


}

#endif
